from calibration_nodes import *
import time
import pandas as pd
import networkx as nx
import numpy
import matplotlib.pyplot as plt

def format_time(time_diff):
    int_time_diff = int(time_diff)
    return '{} hours {} minutes {} seconds'\
        .format(int_time_diff // 3600, int_time_diff % 3600 // 60, int_time_diff % 60)

fault_parameters = -1
uncalibrated = 0
well_calibrated = 1

class Graph:
    def __init__(self):
        '''
        node_status:
            -1 : fault parameters
            0 : uncalibrated
            1 : well calibrated
        '''
        self.source_node = "LoadIon"
        self.current_node = "LoadIon"
        self.current_father_node = None
        self.final_node = "MSGate"
        self.node_number = len(graph_nodes)
        self.node_list = []
        self.node_dict = {}
        self.successor = {}
        self.predecessor_number = {}
        self.pos = None #self.pos用于画图

        ### 用户需要访问的数据
        self.node_status = {}
        self.para_set = {}
        self.last_calibration_time = {}
        ###

        #建立DAG(有向无环图)
        for node in graph_nodes:
            n = node()
            self.node_dict[node.__name__] = n
            self.node_status[node.__name__] = 0
            self.predecessor_number[node.__name__] = 0
            self.successor[node.__name__] = []
            self.last_calibration_time[node.__name__] = 0
            self.para_set[node.__name__] = {}
        for edge in graph_edges:
            self.successor[edge[0].__name__].append(edge[1].__name__)
            self.predecessor_number[edge[1].__name__] += 1

        self.topological_sort()

    def topological_sort(self):
        '''
        Topological sorting for calibration DAG
        Here the 'node_status' only used as a counter
        '''
        l = 0
        self.node_list.append(self.source_node)
        while l < self.node_number:
            current_node = self.node_list[l]
            for next_n in self.successor[current_node]:
                self.node_status[next_n] += 1
                if self.node_status[next_n] == self.predecessor_number[next_n]:
                    self.node_list.append(next_n)
            l += 1
        for node in self.node_list:
            self.node_status[node] = 0

    def spread(self, node, val):
        spread_list = [node]
        while spread_list != []:
            node = spread_list[0]
            for next_n in self.successor[node]:
                self.node_status[next_n] = val
                if not (next_n in spread_list):
                    spread_list.append(next_n)
            spread_list.pop(0)

    def auto_cali_ab_initio(self, starting_node = None):
        if starting_node == None:
            self.node_status[self.source_node] = 0
            self.spread(self.source_node, 0)
        else:
            if type(starting_node) != str:
                starting_node = str(starting_node)
            assert starting_node in self.node_list
            self.node_status[starting_node] = 0
            self.spread(starting_node, 0)
        self.__auto_calibrate()

    def auto_calibrate(self):
        self.__auto_calibrate()

    def __auto_calibrate(self):
        time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
        time_st = time.time()
        print('calibration starts in '+time_str)
        success_flag = True

        for node in self.node_list:
            if self.node_status[node] in [-1, 0]:
                node_object = self.node_dict[node]

                status,para_set = node_object.run()
                self.node_status[node] = status
                if status != 1:
                    print("Calibration failed in node {}, auto_cali terminates.".format(node))
                    success_flag = False
                    self.node_status[node] = -1
                    break
                self.last_calibration_time[node] = time.time()
                self.node_status[node] = status
                if node_object.system_drift(para_set, self.para_set[node]):
                    self.spread(node, 0)
                self.para_set[node] = para_set

        if success_flag:
            print("Auto calibration succeeds!")
            time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            print("Current time: " + time_str)
            time_diff = time.time() - time_st
            print("It takes {} in total."\
                  .format(format_time(time_diff)) )


    def system_time_check(self, starting_node = None):
        uncalibrated_list = []
        if starting_node == None:
            for node in self.node_list:
                node_object = self.node_dict[node]

                status = node_object.check(self.last_calibration_time[node])
                if status < 1:
                    uncalibrated_list.append(node)
        else:
            if type(starting_node) != str:
                starting_node = starting_node.__name__
                assert (starting_node in self.node_list)

                remaning_nodes = [starting_node]
                while remaning_nodes != []:
                    node = remaning_nodes[0]
                    node_object = self.node_dict[node]
                    status = node_object.check(self.last_calibration_time[node])
                    if status == 0:
                        uncalibrated_list.append(node)
                        self.node_status[node] = status

                    for next_n in self.successor[node]:
                        if not (next_n in remaning_nodes):
                            remaning_nodes.append(next_n)
                remaning_nodes.pop(0)

        node_str = ':'
        for idx, node in enumerate(uncalibrated_list):
            node_str += ' ' + node
            if idx < len(uncalibrated_list)-1:
                node_str += ','
            else:
                node_str += '.'
        uncalibrated_number = 0
        for item in self.node_status.items():
            if item[1] < 0:
                uncalibrated_number += 1
        print("System check ends. There are {} nodes to be recalibrated, they are {}"\
              .format(uncalibrated_number, node_str))

    def core_check(self, core_nodes = []):
        assert type(core_nodes) == list
        for node in core_nodes:
            node_object = self.node_dict[node.__name__]
            node = node.__name__

            status, para_set = node_object.run()
            self.node_status[node] = status
            if status == 1:
                if node_object.system_drift(para_set, self.para_set[node]) == True:
                    self.spread(node, 0)
                self.para_set[node] = para_set
            else:
                print("Calibration failed in node {} during core-node-check.".format(node))
                self.spread(node, 0)

    def draw_graph(self):
        dag_dict = {'from':[], 'to': []}
        for node in self.node_list:
            for next_n in self.successor[node]:
                dag_dict['from'].append(node)
                dag_dict['to'].append(next_n)
        dag_dict = pd.DataFrame(dag_dict)
        ### color sheet# -1: 333   0: 000   1: 111
        color_dict = {-1:'#FF0000', 0:'#808080', 1: '#00FF7F'}
        node_color = {'ID':[], 'value':[]}
        for node in self.node_list:
            node_color['ID'].append(node)
            node_color['value'].append(color_dict[self.node_status[node]])
        node_color = pd.DataFrame(node_color)

        G = nx.from_pandas_edgelist(dag_dict, 'from', 'to', create_using = nx.DiGraph())
        node_color = node_color.set_index('ID')
        #print(dag_dict)
        #print(node_color)
        
        if self.pos == None:
            self.pos = nx.spring_layout(G)
        fig = plt.figure( figsize = (10,8))
        nx.draw(G, self.pos, with_labels=True, node_color=numpy.array(node_color['value']), node_size=3000, \
                font_size = 9, alpha=0.9, arrows=True)
        plt.show()

    def reset_pos(self): #如果你对有向图的布局不满意，可以调用此函数重新生成布局
        dag_dict = {'from':[], 'to': []}
        for node in self.node_list:
            for next_n in self.successor[node]:
                dag_dict['from'].append(node)
                dag_dict['to'].append(next_n)
        dag_dict = pd.DataFrame(dag_dict)
        G = nx.from_pandas_edgelist(dag_dict, 'from', 'to', create_using = nx.DiGraph())
        
        self.pos = nx.spring_layout(G)
